package flink.connector.alarm;

import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.apache.flink.table.data.RowData;
import org.apache.flink.util.concurrent.ExecutorThreadFactory;

import com.alibaba.fastjson.JSON;
import flink.bean.WeChatAlarm;
import flink.yarn.YarnClientUtils;
import okhttp3.*;
import org.apache.hadoop.yarn.api.records.ApplicationReport;
import org.apache.hadoop.yarn.exceptions.YarnException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;

import java.io.IOException;
import java.sql.Timestamp;
import java.util.*;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 自定义 {@link WechatAlarmSink} 发送报警到企业微信
 *
 * @author jasonlee
 */
public class WechatAlarmSink extends RichSinkFunction<RowData> {
    private static final Logger LOG = LoggerFactory.getLogger(WechatAlarmSink.class);
    private final String url;
    private final String msgType;
    private final @Nullable List<String> mentionedList = new ArrayList<>();
    private OkHttpClient httpClient;
    private List<WeChatAlarm> weChatAlarmList;
    private Map<String, ApplicationReport> applicationReportMap = new HashMap<>(200);
    private static final int BATCH_SIZE = 10;
    private static final String SPLIT_SEPARATOR = ",";

    public WechatAlarmSink(String url, String msgType, @Nullable String mentionedList) {
        this.url = url;
        this.msgType = msgType;
        if (mentionedList != null) {
            this.mentionedList.addAll(Arrays.asList(mentionedList.split(SPLIT_SEPARATOR)));
        }
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        this.httpClient = new OkHttpClient().newBuilder().build();
        this.weChatAlarmList = new ArrayList<>(BATCH_SIZE);
        ScheduledExecutorService scheduler =
                Executors.newScheduledThreadPool(
                        1, new ExecutorThreadFactory("alarm-connector-pool"));
        scheduler.scheduleWithFixedDelay(
                () -> {
                    synchronized (this) {
                        if (this.weChatAlarmList.size() > 0) {
                            LOG.info("开始执行定时发送企业微信报警, list: {}", weChatAlarmList);
                            sendAlarmToWeChat(this.weChatAlarmList);
                        }
                    }
                },
                5,
                10,
                TimeUnit.SECONDS);

        // 定期扫描 Yarn 上任务的状态
        ScheduledExecutorService schedulerYarnService =
                Executors.newScheduledThreadPool(
                        1, new ExecutorThreadFactory("yarn-application-pool"));
        schedulerYarnService.scheduleWithFixedDelay(
                () -> {
                    synchronized (this) {
                        try {
                            applicationReportMap = YarnClientUtils.getApplicationReport();
                            LOG.info("定时扫描任务执行完成,yarn 上一共有多少个任务 {}", applicationReportMap.size());
                        } catch (IOException | YarnException e) {
                            throw new RuntimeException(e);
                        }
                    }
                },
                30,
                60 * 2,
                TimeUnit.SECONDS);
    }

    @Override
    public void invoke(RowData value, Context context) {
        final String jobName = String.valueOf(value.getString(4));
        // applicationReport 有可能是空的,在提交 Flink 任务的时候 yarn.application.name 和 pipeline.name
        // 必须设置,且要保持一致
        final ApplicationReport applicationReport = this.applicationReportMap.get(jobName);
        final String trackingUrl =
                applicationReport == null ? "" : applicationReport.getTrackingUrl();
        final WeChatAlarm weChatAlarm =
                new WeChatAlarm(
                        String.valueOf(value.getString(0)),
                        String.valueOf(value.getString(1)),
                        String.valueOf(value.getString(2)),
                        String.valueOf(value.getString(3)),
                        // job name
                        jobName,
                        String.valueOf(value.getString(5)),
                        String.valueOf(value.getString(6)),
                        trackingUrl);
        this.weChatAlarmList.add(weChatAlarm);
        if (this.weChatAlarmList.size() >= BATCH_SIZE) {
            LOG.info("缓存超过了 BATCH_SIZE, 开始发送报警: " + weChatAlarmList);
            sendAlarmToWeChat(this.weChatAlarmList);
        }
    }

    @Override
    public void close() {}

    public String generateAlarmContent(WeChatAlarm weChatAlarm, String msgType) {
        Map<String, Object> paramsMap = new HashMap<>(8);
        Map<String, Object> textMap = new HashMap<>(8);
        switch (msgType.toUpperCase()) {
            case "TEXT":
                paramsMap.put("msgtype", "text");
                textMap.put("content", getTextContent(weChatAlarm));
                // 添加想要@的人
                textMap.put("mentioned_list", this.mentionedList);
                paramsMap.put("text", textMap);
                break;
            case "MARKDOWN":
            default:
                paramsMap.put("msgtype", "markdown");
                textMap.put("content", getMarkdownContent(weChatAlarm));
                paramsMap.put("markdown", textMap);
                break;
        }
        return JSON.toJSONString(paramsMap);
    }

    public String getMarkdownContent(WeChatAlarm weChatAlarm) {
        String template =
                "# Flink 任务异常报警请相关开发人员注意。\n"
                        + ">告警级别: <font color=\\\"comment\\\">%s</font>\n"
                        + ">告警类型: <font color=\\\"comment\\\">%s</font>\n"
                        + ">告警详情: <font color=\\\"comment\\\">%s</font>\n"
                        + ">告警时间: <font color=\\\"comment\\\">%s</font>\n"
                        + ">JOB_NAME: <font color=\\\"comment\\\">%s</font>\n"
                        + ">JOB_ID: <font color=\\\"comment\\\">%s</font>\n"
                        + ">METRIC_NAME: <font color=\\\"comment\\\">%s</font>\n"
                        + ">YARN_URL: <font color=\\\"comment\\\">[YARN_URL](%s)</font>\n";
        return combineAlarmContent(template, weChatAlarm);
    }

    public String getTextContent(WeChatAlarm weChatAlarm) {
        String template =
                "Flink 任务异常报警请相关开发人员注意.\n"
                        + "告警级别: %s\"\n"
                        + "告警类型: %s\"\n"
                        + "告警详情: %s\"\n"
                        + "告警时间: %s\"\n"
                        + "JOB_NAME: %s\"\n"
                        + "JOB_ID: %s\"\n"
                        + "METRIC_NAME: %s"
                        + "YARN_URL: %s";
        return combineAlarmContent(template, weChatAlarm);
    }

    public String combineAlarmContent(String template, WeChatAlarm weChatAlarm) {
        return String.format(
                template,
                weChatAlarm.getLevel(),
                weChatAlarm.getType(),
                weChatAlarm.getDetails(),
                weChatAlarm.getAlarmTime(),
                weChatAlarm.getJobName(),
                weChatAlarm.getJobId(),
                weChatAlarm.getMetricName(),
                weChatAlarm.getYarnUrl());
    }

    public void sendAlarmToWeChat(List<WeChatAlarm> weChatAlarmList) {
        LOG.info("当前的线程是: " + Thread.currentThread().getName());
        groupingGetNewest(weChatAlarmList)
                .forEach(
                        weChatAlarm -> {
                            try {
                                sendAlert(weChatAlarm);
                            } catch (IOException e) {
                                throw new RuntimeException(e);
                            }
                        });
        // 发送完报警后要清空 list
        this.weChatAlarmList.clear();
    }

    public void sendAlert(WeChatAlarm weChatAlarm) throws IOException {
        MediaType mediaType = MediaType.parse("application/json");
        RequestBody body =
                RequestBody.create(mediaType, generateAlarmContent(weChatAlarm, this.msgType));
        Request request =
                new Request.Builder()
                        .url(this.url)
                        .addHeader("content-type", "application/json")
                        .post(body)
                        .build();
        Response response = httpClient.newCall(request).execute();
        String result = response.body().string();
        if (response.isSuccessful()) {
            LOG.info("报警发送成功了,接口返回的结果是: {}", result);
        } else {
            LOG.error("报警发送失败了,接口返回的结果是: {}", result);
        }
    }

    /**
     * 对 list 数据根据 jobId,type 分组排序 根据报警时间取最新的一条
     *
     * @param weChatAlarmList 想要发送告警的集合
     * @return 返回每组最新的一条数据
     */
    public List<WeChatAlarm> groupingGetNewest(List<WeChatAlarm> weChatAlarmList) {
        final Map<String, WeChatAlarm> collect =
                weChatAlarmList.stream()
                        .collect(
                                Collectors.toMap(
                                        weChatAlarm ->
                                                weChatAlarm.getJobId() + weChatAlarm.getType(),
                                        Function.identity(),
                                        (a, b) ->
                                                Timestamp.valueOf(a.getAlarmTime()).getTime()
                                                                > Timestamp.valueOf(
                                                                                b.getAlarmTime())
                                                                        .getTime()
                                                        ? a
                                                        : b));
        return new ArrayList<>(collect.values());
    }
}
